// "Copyright [2021] <Copyright Zhancheng Liang, Shengkai Lin>"

#include <vector>
#include "gpu.h"

namespace ai_sim {

Define_Module(Gpu);

void Gpu::initialize() {
    my_global_index = par("my_global_index");
    EV<< "GPU " << my_global_index <<" init finished" << std::endl;
    my_index = par("my_index");
    numNvlink_port = par("numNvlink_port");
    omnetpp::cModule * cpu_module =
        getParentModule()->getSubmodule("cpu", my_index + 1);
    my_cpu = dynamic_cast<Cpu *> (cpu_module);
    ASSERT(my_cpu->getIndex() == getIndex() + 1);
}

void Gpu::handleSend(Flow alloc_flow) {
    // error("stop");
    Enter_Method("handleSend(Flow alloc_flow)");
    FlowMsg *d_msg = new FlowMsg;
    d_msg->setKind(DATA_MSG);
    d_msg->setThe_flow(alloc_flow);
    d_msg->setByteLength(alloc_flow.data_size);
    // error("outport%d", alloc_flow.outportId);
    EV << "Sending a msg with flow src " << alloc_flow.srcIdx <<
        " dest " << alloc_flow.destIdx <<" outport " <<
        gate(alloc_flow.outportId)->getFullPath() << std::endl;
    handleSend(d_msg);
}

void Gpu::handleSend(omnetpp::cMessage * msg) {
    FlowMsg *f_msg = dynamic_cast<FlowMsg *> (msg);
    if (!f_msg) error("Only FlowMsg is allowed here!");
    int output_port = f_msg->getThe_flow().outportId;
    EV << "outport" << output_port << endl;
    omnetpp::cChannel *tmp_channel = gate(output_port)
        ->getTransmissionChannel();
    std::queue<omnetpp::cMessage*> &tmp_queue = msg_queues[output_port];

    if (f_msg->getKind() == DATA_MSG) {
        if (tmp_channel->isBusy()) {
            tmp_queue.push(f_msg);
        } else if (!tmp_queue.empty()) {
            tmp_queue.push(f_msg);
        } else {
            FlowMsg *self_msg = f_msg->dup();
            self_msg->setKind(SELF_MSG);
            send(f_msg, output_port);
            emit(inet::packetSentSignal, f_msg);
            scheduleAt(tmp_channel->getTransmissionFinishTime(),
                self_msg);
        }
    } else if (f_msg->getKind() == SELF_MSG) {
        delete f_msg;
        if (tmp_queue.empty()) {
            return;
        } else {
            FlowMsg* f_msg = dynamic_cast<FlowMsg*> (tmp_queue.front());
            FlowMsg *self_msg = f_msg->dup();
            self_msg->setKind(SELF_MSG);
            send(f_msg, output_port);
            omnetpp::cChannel *tmp_channel = gate(output_port)->getTransmissionChannel();
            tmp_queue.pop();
            scheduleAt(tmp_channel->getTransmissionFinishTime(), self_msg);
        }
    } else {
        error("ERROR TYPE OF MSG %d", msg->getKind());
    }
}

void Gpu::handleMessage(omnetpp::cMessage *msg) {
    if (msg->getKind() == SELF_MSG) {
        handleSend(msg);
    // } else if (msg->getKind() == FLOW_ALLOC_MSG) {
    //     EV << "Send out from gpu "
    //        << my_global_index << endl;
    //     FlowVectorMsg *fa_msg = dynamic_cast<FlowVectorMsg *> (msg);
    //     std::vector<Flow> the_flows = fa_msg->getThe_flows();
    //     delete fa_msg;
    //     numComm_peer = the_flows.size();
    //     for (Flow now_flow : the_flows) {
    //         handleSend(now_flow);
    //     }
    } 
    else if (msg->getKind() == DATA_MSG) 
    {
        EV_INFO << msg->getArrivalGate()->getName()<< std::endl;
        FlowMsg *f_msg = dynamic_cast<FlowMsg *> (msg);
        EV_INFO << "recieve a data msg from gpu "
            << f_msg->getThe_flow().srcIdx
            << " on gpu "<< my_global_index<< std::endl;
        if (f_msg->getThe_flow().destIdx != my_global_index) {
            error("error flow arrived");
        }
        // EV << "Flow id is " << f_msg->getThe_flow().flowId <<
        //     "Downstream flow id is " <<f_msg->getThe_flow().downstream_flows[0] << endl;
        my_cpu->handleReceive(f_msg);
        // f_msg->setByteLength(0);
        // f_msg->setKind(RECEIVED_MSG);
        // send(f_msg, "manager_port$o");
        emit(inet::packetReceivedSignal, f_msg);

        delete f_msg;
    } else {
        delete msg;
        error("UNKNOWN TYPE OF MSG ARRIVED ON GPU!");
    }
}
}  // namespace omnetpp
